{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Edit the below cell's `input_log` variable to the path to your BIN (dataflash) log.\n",
    "The BIN log parser currently assumes the first rangefinder is downward facing."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "from dataclasses import dataclass\n",
    "from collections import namedtuple\n",
    "from datetime import datetime\n",
    "import math\n",
    "import pandas as pd\n",
    "from pathlib import Path\n",
    "from pymavlink import mavutil\n",
    "from pyproj import Geod\n",
    "from scipy.spatial.transform import Rotation as R\n",
    "\n",
    "\n",
    "input_log = Path(\"00000125.BIN\")\n",
    "stem = input_log.stem  # '00000125' from '00000125.BIN'\n",
    "output_csv = Path(f\"projected_{stem}.csv\")\n",
    "\n",
    "# Open .bin log\n",
    "log = mavutil.mavlink_connection(str(input_log))\n",
    "\n",
    "# Create a geodesic calculator\n",
    "geod = Geod(ellps='WGS84')\n",
    "\n",
    "# Store messages as lists of tuples\n",
    "attitudes = []\n",
    "ahrs_positions = []\n",
    "rangefinder = []\n",
    "terrain_data = []\n",
    "\n",
    "\n",
    "Ahrs = namedtuple(\"Ahrs\", [\"timestamp\", \"lat\", \"lon\", \"alt\", \"roll\", \"pitch\", \"yaw\"])\n",
    "Rangefinder = namedtuple(\"Rangefinder\", [\"timestamp\", \"distance\"])\n",
    "Terrain = namedtuple(\"Terrain\", [\"timestamp\", \"terr_height\"])\n",
    "\n",
    "@dataclass\n",
    "class GroundPoint:\n",
    "    timestamp: float\n",
    "    rng_lat: float\n",
    "    rng_lon: float\n",
    "    rng_elev: float\n",
    "    ahrs_lat: float\n",
    "    ahrs_lon: float\n",
    "    ahrs_alt: float\n",
    "    distance: float\n",
    "    terr_height: float\n",
    "    \n",
    "\n",
    "while True:\n",
    "    msg = log.recv_match(blocking=False)\n",
    "    if msg is None:\n",
    "        break\n",
    "\n",
    "    t = datetime.fromtimestamp(msg._timestamp)\n",
    "    mtype = msg.get_type()\n",
    "\n",
    "    if msg.get_type() == \"AHR2\":\n",
    "        # Lat and Lon in degrees, alt in m\n",
    "        ahrs_positions.append(Ahrs(\n",
    "            timestamp=t,\n",
    "            lat=msg.Lat,\n",
    "            lon=msg.Lng,\n",
    "            alt=msg.Alt,\n",
    "            roll=math.radians(msg.Roll),\n",
    "            pitch=math.radians(msg.Pitch),\n",
    "            yaw=math.radians(msg.Yaw),\n",
    "        ))\n",
    "    elif msg.get_type() == \"RFND\" and msg.Stat == 4: # Only use healthy readings, see AP_Rangefinder::Status::Good\n",
    "        rangefinder.append(Rangefinder(\n",
    "            timestamp=t,\n",
    "            distance=msg.Dist\n",
    "        ))\n",
    "    elif msg.get_type() == \"TERR\":\n",
    "        terrain_data.append(Terrain(\n",
    "            timestamp=t,\n",
    "            terr_height=msg.TerrH\n",
    "        ))\n",
    "\n",
    "ahrs_df = pd.DataFrame(ahrs_positions).set_index(\"timestamp\").sort_index()\n",
    "rfnd_df = pd.DataFrame(rangefinder).set_index(\"timestamp\").sort_index()\n",
    "terr_df = pd.DataFrame(terrain_data).set_index(\"timestamp\").sort_index()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Compute estimated ground points\n",
    "ground_points: list[GroundPoint] = []\n",
    "\n",
    "assert ahrs_positions, \"No AHR2 messages found\"\n",
    "assert rangefinder, \"No valid RFND messages found\"\n",
    "assert terrain_data, \"No TERR messages found\"\n",
    "\n",
    "for t, row in rfnd_df.iterrows():\n",
    "    dist = row.distance\n",
    "\n",
    "    # Efficient nearest neighbor lookup with .asof() for datetime index\n",
    "    try:\n",
    "        ahrs_row = ahrs_df.loc[ahrs_df.index.asof(t)]\n",
    "        terr_row = terr_df.loc[terr_df.index.asof(t)]\n",
    "    except KeyError:\n",
    "        continue  # skip if no data available at or before timestamp\n",
    "\n",
    "    # Unpack required fields\n",
    "    lat, lon, alt, roll, pitch, yaw  = ahrs_row.lat, ahrs_row.lon, ahrs_row.alt, ahrs_row.roll, ahrs_row.pitch, ahrs_row.yaw\n",
    "    terr_height = terr_row.terr_height\n",
    "\n",
    "    # Build the rotation using Euler angles (in radians)\n",
    "    rotation = R.from_euler('xyz', [roll, pitch, yaw])\n",
    "\n",
    "    # Rotate the down vector (in body frame)\n",
    "    ned_vector = rotation.apply([0, 0, dist])\n",
    "\n",
    "    # Calculate ground point in NED frame\n",
    "    north_offset = ned_vector[0]\n",
    "    east_offset = ned_vector[1]\n",
    "    down_offset = ned_vector[2]\n",
    "\n",
    "    ground_alt = alt - down_offset  # down_offset is positive, so subtract it to get estimated ground height\n",
    "\n",
    "    # Project horizontal lat/lon offset\n",
    "    azimuth = math.degrees(math.atan2(east_offset, north_offset))\n",
    "    horizontal_distance = math.hypot(north_offset, east_offset)\n",
    "    g_lon, g_lat, _ = geod.fwd(lon, lat, azimuth, horizontal_distance)\n",
    "\n",
    "    ground_points.append(GroundPoint(\n",
    "        timestamp=t,\n",
    "        rng_lat=g_lat,\n",
    "        rng_lon=g_lon,\n",
    "        rng_elev=ground_alt,\n",
    "        ahrs_lat=lat,\n",
    "        ahrs_lon=lon,\n",
    "        ahrs_alt=alt,\n",
    "        distance=dist,\n",
    "        terr_height=terr_height\n",
    "    ))\n",
    "\n",
    "if not ground_points:\n",
    "    print(\"Error: No projected ground point data found.\")\n",
    "    exit(1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Convert to DataFrame\n",
    "\n",
    "import pandas as pd\n",
    "\n",
    "df = pd.DataFrame([gp.__dict__ for gp in ground_points])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Get the latitude/longitude extent for input to OpenTopography\n",
    "\n",
    "min_lat = df[\"rng_lat\"].min()\n",
    "max_lat = df[\"rng_lat\"].max()\n",
    "min_lon = df[\"rng_lon\"].min()\n",
    "max_lon = df[\"rng_lon\"].max()\n",
    "\n",
    "assert min_lat < max_lat\n",
    "assert min_lon < max_lon\n",
    "\n",
    "print(f\"Latitude range: ({min_lat},{max_lat})\")\n",
    "print(f\"Longitude range: ({min_lon},{max_lon})\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To compare to other elevation data sources, these must be manually downloaded from OpenTopography.\n",
    "Once extracted into this directory, rename them and update paths in the below block.\n",
    "Add as many datasets for comparison as you like. GDAL supports many different file formats."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "# Load elevation from GeoTIFF using GDAL\n",
    "\n",
    "from osgeo import gdal\n",
    "\n",
    "dataset_paths = [Path(\"cop30.tif\"), Path(\"AW3D30.tif\")]\n",
    "for dataset_path in dataset_paths:\n",
    "    dataset = gdal.Open(str(dataset_path))\n",
    "    assert dataset is not None, f\"Dataset '{dataset_path}' failed to load\"\n",
    "    band = dataset.GetRasterBand(1)\n",
    "    transform = dataset.GetGeoTransform()\n",
    "\n",
    "    # Convert lat/lon to row/col\n",
    "    def latlon_to_rowcol(lat, lon):\n",
    "        inv_transform = gdal.InvGeoTransform(transform)\n",
    "        px, py = gdal.ApplyGeoTransform(inv_transform, lon, lat)\n",
    "        return int(py), int(px)\n",
    "\n",
    "    elevations = []\n",
    "    for _, row in df.iterrows():\n",
    "        r, c = latlon_to_rowcol(row.rng_lat, row.rng_lon)\n",
    "        grid = band.ReadAsArray(c, r, 1, 1)\n",
    "        if grid is not None:\n",
    "            elevation = grid[0][0]\n",
    "        else:\n",
    "            # Out of range of the dataset, common with LOG_REPLAY before the EKF initializes\n",
    "            elevation = math.nan\n",
    "        elevations.append(elevation)\n",
    "\n",
    "    df[dataset_path.stem + \"_elevation\"] = elevations"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# plot using plotly, good for smaller logs (<30 minutes)\n",
    "\n",
    "import plotly.graph_objs as go\n",
    "\n",
    "fig = go.Figure()\n",
    "for col in df.columns:\n",
    "    stems = [f\"{p.stem}_elevation\" for p in dataset_paths]\n",
    "    cols = (\"ahrs_alt\", \"rng_elev\", \"terr_height\", *stems)\n",
    "    if col in cols:\n",
    "        print(f\"Adding {col}\")\n",
    "        y = pd.to_numeric(df[col], errors='coerce')\n",
    "        \n",
    "        fig.add_trace(go.Scattergl(x=df[\"timestamp\"], y=y, mode='lines+markers', name=col))\n",
    "\n",
    "fig.update_layout(title=f\"Terrain Comparison Over Time for {stem}\", xaxis_title=\"Time\", yaxis_title=\"Altitude (m)\", height=600)\n",
    "fig.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Save plot\n",
    "\n",
    "import plotly.io as pio\n",
    "pio.write_html(fig, f'terrain_pyplot_{stem}.html')"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "oi_kernel",
   "language": "python",
   "name": "oi_kernel"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
